import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.colors import Normalize
import matplotlib.colors as colors
from scipy.interpolate import interpn
import matplotlib.ticker as ticker


def colorbar_fmt(x, pos):
    return r'10^{{{}}}$'.format(x)


def filter_densitiy(data, x_e, y_e, threashold):
    for i in range(len(data)):
        for j in range(len(data[i])):
            if data[i][j] > threashold:
                print(x_e[i], y_e[j], x_e[i+1], y_e[i+1])


def scatter_density(z):
    ax = plt.subplot(122)
    x = np.arange(0, len(z), 1)
    ax.scatter(x, z)


def density_scatter(x, y, ticks, tick_labels, ctransform, x_ticks, y_ticks, ax=None, sort=True, bins=20, cmap='jet' ,**kwargs):
    """
    Scatter plot colored by 2d histogram
    """
    if ax is None:
        ax = plt.subplot()
    data, x_e, y_e = np.histogram2d(x, y, bins=bins, density=True)
    filter_densitiy(data, x_e, y_e, 100)
    # z: 每个点的颜色
    z = interpn((0.5*(x_e[1:] + x_e[:-1]), 0.5*(y_e[1:]+y_e[:-1])),
                data, np.vstack([x, y]).T, method="linear", bounds_error=False)

    # To be sure to plot all data
    z[np.where(np.isnan(z))] = 0.0

    # Sort the points by density, so that the densest points are plotted last
    if sort:
        idx = z.argsort()
        x, y, z = x[idx], y[idx], z[idx]

    if ctransform:
        print("before transform, zmin:", np.min(z), ",zmax:", np.max(z))
        z = ctransform(z)
        print("after transform, zmin:", np.min(z), ",zmax:", np.max(z))
        ticks = ctransform(ticks)

    # density = z
    # print("zmin:", np.min(z), ",zmax:", np.max(z))
    # norm = Normalize(vmin=np.min(z), vmax=np.max(z))
    # z = norm(z) * 100 + 1
    # print("zmin:", np.min(z), ",zmax:", np.max(z))

    # z[np.where(np.isinf(z))] = 0.0
    scatter = ax.scatter(x, y, c=z, cmap=cmap, **kwargs)
    cbar = plt.figure(1).colorbar(mappable=scatter,
                                  ax=ax, format=ticker.FuncFormatter(colorbar_fmt), ticks=ticks)
    cbar.ax.set_yticklabels(tick_labels)
    cbar.ax.set_ylabel('Density')
    if x_ticks is not None:
        ax.set_xticks(x_ticks)
    if y_ticks is not None:
        ax.set_yticks(y_ticks)
    plt.show()
    return ax


def draw_heatmap(file, bins=[50, 50], ticks=[0.1, 1, 10, 100, 1000], tick_labels=[r'$10^{-1}$', r'$10^0$', r'$10^1$', r'$10^2$', r'$10^3$'], ctransform=None, x_ticks=None, y_ticks=None):
    geos = []
    with open(file) as csv:
        for line in csv.readlines():
            tokens = line.split(",")
            if y_ticks and len(y_ticks):
                if y_ticks[0] < float(tokens[0]) < y_ticks[-1]:
                    geos.append((tokens[0], tokens[1]))
            else:
                geos.append((tokens[0], tokens[1]))
    lats = np.array([np.float64(g[0]) for g in geos])
    lons = np.array([np.float64(g[1]) for g in geos])
    density_scatter(lons, lats, bins=bins, s=8, ticks=ticks,
                    tick_labels=tick_labels,
                    ctransform=ctransform,
                    x_ticks=x_ticks,
                    y_ticks=y_ticks)


def sg_density_transform(z):
    # norm = Normalize(vmin=np.min(z), vmax=np.max(z))
    # z = norm(z) * 100 + 1
    z = np.log10(z)
    return z


if __name__ == "__main__":
    # draw_heatmap("checkin_sg.csv", ctransform=sg_density_transform,
    #              x_ticks=np.arange(103.65, 104, 0.1),
    #              y_ticks=np.arange(1.2, 1.45, 0.05))
    # draw_heatmap("checkin_ld.csv", ctransform=sg_density_transform,)
    draw_heatmap("data.csv", ctransform=sg_density_transform)
    # test_hist2d()
